from google.protobuf.descriptor_pb2 import (
    DescriptorProto as MessageDP,
    EnumDescriptorProto as EnumDP,
    FieldDescriptorProto as FieldDP,
    FileDescriptorProto as FileDP,
    ServiceDescriptorProto as ServiceDP,
    OneofDescriptorProto as OneofDP,
)
from google.protobuf.descriptor_pool import DescriptorPool

from .descriptor_ext import ProtoFile
from .options import Options
from .utils import (
    get_field_name,
    get_ts_type,
    get_ts_oneof_type,
    is_message,
    split_fields,
    from_napi_value,
    to_napi_value,
    collect_oneof_fields,
    to_json_name,
)


class TemplateRenderer:
    def __init__(self, pool: DescriptorPool, opts: Options) -> None:
        self.pool = pool
        self.opts = opts
        self.classes = []
        self.enums = []

    def clear(self):
        self.classes = []
        self.enums = []

    def get_type_name_in_cc(self, name: str, parent_names: list[str]) -> str:
        return f"{'_'.join(parent_names)}_{name}" if parent_names else name

    def get_type_name_in_ts(self, name: str, parent_names: list[str]) -> str:
        return (
            f"{'_'.join(parent_names)}_{name}"
            if parent_names and not self.opts.enable_namespace
            else name
        )

    ################################## .napi.h ###################################

    # ---------- 枚举 ----------
    def enum_stub_header(
        self, enum: EnumDP, parent_names: list[str] = []
    ) -> list[str]:
        enum_name = self.get_type_name_in_cc(enum.name, parent_names)
        self.enums.append(enum_name)
        if self.opts.enable_namespace:
            return [f"napi_value InitEnum_{enum_name}(napi_env env);"]
        return [f"void InitEnum_{enum_name}(napi_env env, napi_value exports);"]

    # ---------- 消息 ----------
    def message_stub_header(
        self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = ""
    ) -> list[str]:
        # Skip map entry types - they don't need stub classes
        if msg.options and msg.options.map_entry:
            return []

        message_name = self.get_type_name_in_cc(msg.name, parent_names)
        self.classes.append(message_name)

        fields = [
            f"""\
    static napi_value Get{field.name.capitalize()}(napi_env env, napi_callback_info info);
    static napi_value Set{field.name.capitalize()}(napi_env env, napi_callback_info info);
"""
            for field in msg.field
            if not field.HasField("oneof_index")
        ]
        oneofs = [
            f"""\
    static napi_value Get{field.name.capitalize()}(napi_env env, napi_callback_info info);
    static napi_value Set{field.name.capitalize()}(napi_env env, napi_callback_info info);
    static napi_value _Set{field.name.capitalize()}(napi_env env, napi_value arg, {message_name} *impl);
"""
            for field in msg.oneof_decl
        ]
        namespace = pkg_name.replace(".", "::")
        return [
            f"""\
class {message_name}_Stub {{
public:
    {message_name}_Stub() = delete;
    ~{message_name}_Stub() = default;

    static napi_value Init(napi_env env, napi_value exports);
    static void Finalizer(napi_env env, void *data, void *hint) {{
        if (NapiUtils::g_ENABLE_ARENA) {{
            auto impl = reinterpret_cast<{namespace}::{message_name} *>(data);;
            NapiUtils::DecArenaRefCount(impl->GetArena());
        }} else {{
            delete reinterpret_cast<{namespace}::{message_name} *>(data);
        }}
    }}
    // If a C++ pointer is wrapped in @p nv, use the pointer to initialize @p impl.
    // Otherwise, try using the properties of @p nv.
    static void InitByNapiObj(napi_env env, napi_value nv, {namespace}::{message_name} *impl);

    static thread_local napi_ref ctorRef;

private:
    static napi_value New(napi_env env, napi_callback_info info);
    {"static napi_value Clone(napi_env env, napi_callback_info info);" if self.opts.enable_clone else ""}
    static napi_value ToBinary(napi_env env, napi_callback_info info);
    static napi_value FromBinary(napi_env env, napi_callback_info info);
""",
            *fields,
            *oneofs,
            f"}}; // class {message_name}_Stub",
        ]

    def one_init_per_header(self, proto: ProtoFile):
        enum_inites = (
            [f"    InitEnum_{enum}(env, exports);" for enum in self.enums] if not self.opts.enable_namespace else []
        )
        class_inites = [f"    {cls}_Stub::Init(env, exports);" for cls in self.classes]
        return [
            f"""\
inline napi_value Init_{proto.file_basename}(napi_env env, napi_value exports) {{
    {f'using namespace {proto.package_name.replace(".", "::")};' if proto.package_name else ""}
""",
            *enum_inites,
            *class_inites,
            """\
    return exports;
}
""",
        ]

    def all_in_one_init(self, inites: list[tuple[str, ProtoFile]]):
        includes = [f'#include "{init[0]}"' for init in inites]
        calls = [f"    Init_{init[1].file_basename}(env, exports);" for init in inites]
        codes = [
            *includes,
            "",
            f"""\
inline napi_value InitAllInOne(napi_env env, napi_value exports) {{
    napi_property_descriptor desc[] = {{
        {('' if self.opts.enable_arena_switch else '// ') + '{ "setArenaEnabled", nullptr, NapiUtils::SetArenaEnabled, nullptr, nullptr, nullptr, napi_default, nullptr }'},
        {('' if self.opts.enable_json else '// ') + '{ "toJsonString", nullptr, NapiUtils::ToJsonString, nullptr, nullptr, nullptr, napi_default, nullptr }'},
    }};
    napi_define_properties(env, exports, sizeof(desc) / sizeof(desc[0]), desc);\
""",
            *calls,
            """\
    return exports;
}
""",
        ]
        return "\n".join(codes)

    ################################## .napi.cc ###################################

    # ---------- 枚举 ----------
    def enum_stub(self, enum: EnumDP, parent_names: list[str] = [], pkg_name: str = "") -> list[str]:
        enum_name_in_cc = self.get_type_name_in_cc(enum.name, parent_names)
        enum_name_in_ts = self.get_type_name_in_ts(enum.name, parent_names)
        fields_1 = [
            f'        {{ "{field.number}", nullptr, nullptr, nullptr, nullptr, NapiUtils::CreateScalarValue<std::string>(env, "{field.name}"), napi_enumerable, nullptr}},'
            for i, field in enumerate(enum.value)
        ]
        fields_2 = [
            f'        {{ "{field.name}", nullptr, nullptr, nullptr, nullptr, NapiUtils::CreateScalarValue<int32_t>(env, {field.number}), napi_enumerable, nullptr}},'
            for i, field in enumerate(enum.value)
        ]
        func_signature = (
            f"napi_value InitEnum_{enum_name_in_cc}(napi_env env)"
            if self.opts.enable_namespace
            else f"void InitEnum_{enum_name_in_cc}(napi_env env, napi_value exports)"
        )
        return [
            f"""
{func_signature} {{
    const napi_property_descriptor desc[] = {{\
""",
            *fields_1,
            *fields_2,
            f"""\
    }};
    napi_value enumObj = nullptr;
    napi_create_object(env, &enumObj);
    napi_define_properties(env, enumObj, sizeof(desc) / sizeof(desc[0]), desc);
    {"return enumObj;" if self.opts.enable_namespace else f'napi_set_named_property(env, exports, "{enum_name_in_ts}", enumObj);'}
}}
""",
        ]

    def message_stub(self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = "") -> list[str]:
        message_name = self.get_type_name_in_cc(msg.name, parent_names)

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        return [f"thread_local napi_ref {message_name}_Stub::ctorRef = nullptr;"]

    def message_init_stub(
        self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = ""
    ) -> list[str]:
        message_name_in_cc = self.get_type_name_in_cc(msg.name, parent_names)
        message_name_in_ts = self.get_type_name_in_ts(msg.name, parent_names)

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        # 将枚举作为静态成员变量设置入类中
        static_enum_fields = [
            f'        {{"{enum.name}", nullptr, nullptr, nullptr, nullptr, InitEnum_{self.get_type_name_in_cc(enum.name, [*parent_names, message_name_in_cc])}(env), napi_static, nullptr}},'
            for enum in msg.enum_type
        ] if self.opts.enable_namespace else []

        fields = [
            f'        {{"{field.json_name}", nullptr, nullptr, Get{field.name.capitalize()}, Set{field.name.capitalize()}, nullptr, napi_default, nullptr}},'
            for field in msg.field
            if not field.HasField("oneof_index")
        ]
        fields.extend(
            [
                f'        {{"{to_json_name(field.name)}", nullptr, nullptr, Get{field.name.capitalize()}, Set{field.name.capitalize()}, nullptr, napi_default, nullptr}},'
                for field in msg.oneof_decl
            ]
        )
        outer_namespace_str = f"\"{'/'.join([*pkg_name.split('.'), *parent_names])}\""
        outer_namespace = (
            f"NapiUtils::GetNamespaceObj(env, exports, {outer_namespace_str})"
            if self.opts.enable_namespace and (pkg_name or parent_names)
            else "exports"
        )
        return [
            f"""\
napi_value {message_name_in_cc}_Stub::Init(napi_env env, napi_value exports) {{
    const napi_property_descriptor properties[] = {{
""",
            *static_enum_fields,
            " " * 8 + "// all fields",
            *fields,
            f"""\
        // methods
        {'{"clone", nullptr, Clone, nullptr, nullptr, nullptr, napi_default, nullptr},' if self.opts.enable_clone else ''}
        {{"toBinary", nullptr, ToBinary, nullptr, nullptr, nullptr, napi_default, nullptr}},
        {{"fromBinary", nullptr, FromBinary, nullptr, nullptr, nullptr, napi_static, nullptr}},
    }};
    napi_value ctor = nullptr;
    napi_define_class(env, "{message_name_in_ts}", NAPI_AUTO_LENGTH, New, nullptr, sizeof(properties) / sizeof(properties[0]),
                    properties, &ctor);
    napi_create_reference(env, ctor, 1, &ctorRef);
    napi_set_named_property(env, {outer_namespace}, "{message_name_in_ts}", ctor);\
""",
            f"""
    NapiUtils::SetNamespaceObj({outer_namespace_str}, "{message_name_in_ts}", ctor);
""" if static_enum_fields else "",
"""\
    return exports;
}
"""
        ]

    def message_init_by_napi_obj_stub(
        self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = ""
    ) -> list[str]:
        message_name = self.get_type_name_in_cc(msg.name, parent_names)

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        required_fields, optional_fields = split_fields(msg)
        oneofs = collect_oneof_fields(msg)
        required_fields = [
            f"""\
    if (auto prop = NapiUtils::GetProperty(env, nv, "{field.json_name}")) {{
        {from_napi_value(field, f"*prop", pkg_name, self.pool)}
    }}\
"""
            for field in required_fields
            if not field.HasField("oneof_index")
        ]
        oneof_fields = [
            f"""\
    if (auto prop = NapiUtils::GetProperty(env, nv, "{to_json_name(oneof[0].name)}")) {{
        _Set{oneof[0].name.capitalize()}(env, *prop, impl);
    }}\
"""
            for oneof in oneofs
        ]
        optional_fields = [
            f"""\
    if (auto prop = NapiUtils::GetProperty(env, nv, "{field.json_name}")) {{
        {from_napi_value(field, f"*prop", pkg_name, self.pool)}
    }}\
"""
            for field in optional_fields
        ]
        full_type_name = f'{pkg_name.replace(".", "::")}::{message_name}'
        return [
            f"""\
void {message_name}_Stub::InitByNapiObj(napi_env env, napi_value nv, {full_type_name} *impl) {{
    if (auto cv = NapiUtils::Unwrap<{full_type_name}>(env, nv, false)) {{
        impl->CopyFrom(*cv);
        return;
    }}
""",
            *required_fields,
            *oneof_fields,
            *optional_fields,
            "}",
        ]

    def message_ctor_stub(self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = "") -> list[str]:
        """
        生成消息的构造函数
        """
        message_name = self.get_type_name_in_cc(msg.name, parent_names)

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        full_type_name = f'{pkg_name.replace(".", "::")}::{message_name}'
        return [
            f"""\
napi_value {message_name}_Stub::New(napi_env env, napi_callback_info info) {{
    auto [thisArg, args] = NapiUtils::GetThisAndArgs(env, info, 1);
    if (args.size() == 1 && NapiUtils::IsNull(env, args[0])) {{
        return thisArg;
    }}

    {full_type_name} *impl = nullptr;
    if (NapiUtils::g_ENABLE_ARENA) {{
        auto arena = NapiUtils::NewArena();
        impl = google::protobuf::Arena::Create<{full_type_name}>(arena);
    }} else {{
        impl = new {full_type_name}();
    }}
    if (args.size() != 0 && !NapiUtils::IsUndefined(env, args[0])) {{
        {message_name}_Stub::InitByNapiObj(env, args[0], impl);
    }}
    napi_wrap(env, thisArg, impl, Finalizer, nullptr, nullptr);
    return thisArg;
}}
""",
        ]

    def message_to_binary_stub(self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = "") -> list[str]:
        message_name = self.get_type_name_in_cc(msg.name, parent_names)
        namespace = pkg_name.replace(".", "::")

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        return [
            f"""\
napi_value {message_name}_Stub::ToBinary(napi_env env, napi_callback_info info) {{
    return NapiUtils::ToBinary<{namespace}::{message_name}>(env, info);
}}
""",
        ]

    def message_from_binary_stub(
        self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = ""
    ) -> list[str]:
        message_name = self.get_type_name_in_cc(msg.name, parent_names)
        namespace = pkg_name.replace(".", "::")

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        return [
            f"""\
napi_value {message_name}_Stub::FromBinary(napi_env env, napi_callback_info info) {{
    {namespace}::{message_name} *impl = nullptr;
    if (NapiUtils::g_ENABLE_ARENA) {{
        auto arena = NapiUtils::NewArena();
        impl = google::protobuf::Arena::Create<{namespace}::{message_name}>(arena);
    }} else {{
        impl = new {namespace}::{message_name}();
    }}
    return NapiUtils::FromBinary(env, info, impl, ctorRef, Finalizer);
}}
""",
        ]

    def message_clone_stub(self, msg: MessageDP, parent_names: list[str] = [], pkg_name: str = "") -> list[str]:
        message_name = self.get_type_name_in_cc(msg.name, parent_names)
        namespace = pkg_name.replace(".", "::")

        # Skip map entry types
        if msg.options and msg.options.map_entry:
            return []

        return [
            f"""\
napi_value {message_name}_Stub::Clone(napi_env env, napi_callback_info info) {{
    return NapiUtils::Clone<{namespace}::{message_name}, {namespace}::{message_name}_Stub>(env, info, ctorRef, Finalizer);
}}
""",
        ]

    # ---------- 字段 ----------
    def field_stub(self, field: FieldDP, msg_name: str, pkg_name: str) -> list[str]:
        if field.HasField("oneof_index"):
            return []

        return [
            f"""\
napi_value {msg_name}_Stub::Get{field.name.capitalize()}(napi_env env, napi_callback_info info) {{
    auto thisArg = NapiUtils::GetThis(env, info);
    {msg_name} *impl = NapiUtils::Unwrap<{msg_name}>(env, thisArg);
    if (!impl) {{
        return NapiUtils::CreateUndefined(env);
    }}
    {f"if (!impl->has_{get_field_name(field)}()) {{ return NapiUtils::CreateUndefined(env); }}" if is_message(field) else ''}
    return {to_napi_value(field, pkg_name, self.pool)};
}}
""",
            f"""\
napi_value {msg_name}_Stub::Set{field.name.capitalize()}(napi_env env, napi_callback_info info) {{
    auto [thisArg, args] = NapiUtils::GetThisAndArgs(env, info, 1);
    if ({msg_name} *impl = NapiUtils::Unwrap<{msg_name}>(env, thisArg)) {{
        {from_napi_value(field, "args[0]", pkg_name, self.pool)}
    }}
    return NapiUtils::CreateUndefined(env);
}}
""",
        ]

    # ---------- oneof ----------
    def oneof_stub(
        self, oneof: OneofDP, fields: list[FieldDP], msg_name: str, pkg_name: str
    ) -> list[str]:
        capitalize_initial_letter = lambda s: s[0].upper() + s[1:]
        oneof_kind_in_cc = lambda name: "".join([capitalize_initial_letter(part) for part in name.split("_")])
        cases_in_getter = [
            f"""\
    case {msg_name}::k{oneof_kind_in_cc(field.name)}: {{
        case_ = NapiUtils::CreateScalarValue<std::string>(env, "{field.json_name}");
        value = {to_napi_value(field, pkg_name, self.pool)};
        break;
    }}\
"""
            for field in fields
        ]
        cases_in_setter = [
            f"""\
    if (kase == "{field.json_name}") {{
        {from_napi_value(field, "value", pkg_name, self.pool, True)}
        return undefined;
    }}\
"""
            for field in fields
        ]
        return [
            f"""\
napi_value {msg_name}_Stub::Get{oneof.name.capitalize()}(napi_env env, napi_callback_info info) {{
    auto thisArg = NapiUtils::GetThis(env, info);
    {msg_name} *impl = nullptr;
    napi_unwrap(env, thisArg, reinterpret_cast<void **>(&impl));
    napi_value case_ = nullptr;
    napi_value value = nullptr;
    switch (impl->{oneof.name}_case()) {{\
""",
            *cases_in_getter,
            f"""\
    case {msg_name}::{oneof.name.upper()}_NOT_SET:
        case_ = NapiUtils::CreateUndefined(env);
    }}
    napi_value obj = nullptr;
    napi_create_object(env, &obj);
    napi_set_named_property(env, obj, "case", case_);
    napi_set_named_property(env, obj, "value", value);
    return obj;
}}
""",
            f"""\
napi_value {msg_name}_Stub::Set{oneof.name.capitalize()}(napi_env env, napi_callback_info info) {{
    auto [thisArg, args] = NapiUtils::GetThisAndArgs(env, info, 1);
    {msg_name} *impl = nullptr;
    napi_unwrap(env, thisArg, reinterpret_cast<void **>(&impl));

    return _Set{oneof.name.capitalize()}(env, args[0], impl);
}}
""",
            f"""\
napi_value {msg_name}_Stub::_Set{oneof.name.capitalize()}(napi_env env, napi_value arg, {msg_name} *impl) {{
    napi_value case_ = nullptr;
    napi_get_named_property(env, arg, "case", &case_);
    napi_valuetype caseType;
    napi_typeof(env, case_, &caseType);
    if (caseType == napi_undefined) {{
        return NapiUtils::CreateUndefined(env);
    }}
    napi_value value = nullptr;
    napi_get_named_property(env, arg, "value", &value);
    auto kase = NapiUtils::GetScalarValue<std::string>(env, case_);
    napi_value undefined = NapiUtils::CreateUndefined(env);
""",
            *cases_in_setter,
            """\
    return undefined;
}
""",
        ]

    # ---------- 服务 ----------
    def service_stub(self, svc: ServiceDP) -> list[str]:
        return [
            f"class {svc.name}Stub {{",
            "public:",
            f"    {svc.name}Stub() = delete;",
            f"    ~{svc.name}Stub() = default;",
            "    // TODO(user): 生成 RPC 方法存根",
            "};",
        ]

    # ---------- 扩展 ----------
    def extension_stub(self, ext) -> list[str]:
        return [f"// TODO(user): 生成扩展字段 {ext.name} 的存根"]

    # ---------- 文件选项 ----------
    def file_options_stub(self, options) -> list[str]:
        return ["// TODO(user): 根据文件级选项生成代码"]

    ################################## .d.ts ###################################

    # ---------- 枚举 ----------
    def enum_stub_dts(
        self,
        enum: EnumDP,
        parent_name: str = "",
    ) -> list[str]:
        enum_name = f"{parent_name}_{enum.name}" if parent_name else enum.name
        enum_values = [f"  {field.name} = {field.number}," for field in enum.value]
        posible_export = "export " if not self.opts.enable_namespace else ""
        return (
            [f"{posible_export}enum {enum_name} {{", *enum_values, "}"]
            if self.opts.disable_const_enum
            else [f"export const enum {enum_name} {{", *enum_values, "}"]
        )

    # ---------- 消息 ----------
    def message_stub_dts(
        self,
        msg: MessageDP,
        parent_name: str = "",
        pkg_name: str = "",
    ) -> list[str]:
        message_name = f"{parent_name}_{msg.name}" if parent_name else msg.name

        # Skip map entry types for TypeScript definitions
        if msg.options and msg.options.map_entry:
            return []

        possible_null = " | null" if self.opts.enable_namespace else ""
        possible_export = "export " if not self.opts.enable_namespace else ""
        required_fields, optional_fields = split_fields(msg)
        required_members = [
            f"  {field.json_name}!: {get_ts_type(field, pkg_name, self.pool, self.opts.enable_namespace)};" for field in required_fields
        ]
        required_members_in_if = [
            f"  {field.json_name}?: {get_ts_type(field, pkg_name, self.pool, self.opts.enable_namespace)}{possible_null};" for field in required_fields
        ]
        optional_members = [
            f"  {field.json_name}{'?' if is_message(field) else ''}: {get_ts_type(field, pkg_name, self.pool, self.opts.enable_namespace)}{possible_null};"
            for field in optional_fields
        ]
        optional_members_in_if = [
            f"  {field.json_name}?: {get_ts_type(field, pkg_name, self.pool, self.opts.enable_namespace)}{possible_null};"
            for field in optional_fields
        ]
        oneofs = collect_oneof_fields(msg)
        oneof_members = [
            f"  {to_json_name(oneof.name)}: {get_ts_oneof_type(fields, pkg_name, self.pool, self.opts.enable_namespace)};"
            for oneof, fields in oneofs
        ]
        oneof_members_in_if = [
            f"  {to_json_name(oneof.name)}?: {get_ts_oneof_type(fields, pkg_name, self.pool, self.opts.enable_namespace)}{possible_null};"
            for oneof, fields in oneofs
        ]
        prefix = "I" if self.opts.enable_namespace else "$"
        interface = (
            [
                f"{possible_export}interface {prefix}{message_name} {{",
                *required_members_in_if,
                *optional_members_in_if,
                *oneof_members_in_if,
                f"}} // interface {prefix}{message_name}"
            ]
            if self.opts.enable_namespace
            else [f"export interface ${message_name} extends Partial<{message_name}> {{}}"]
        )
        return [
            *interface,
            f"{possible_export}class {message_name} {f'implements {prefix}{message_name} ' if self.opts.enable_namespace else ''}{{",
            "\n".join([*required_members, *optional_members, *oneof_members]),
            f"  constructor(init?: {prefix}{message_name});",
            f"  clone(): {message_name};" if self.opts.enable_clone else "",
            f"  toBinary(): Uint8Array;",
            f"  static fromBinary(data: Uint8Array): {message_name};",
            "}",
        ]

    def index_ts(self, symbol_map: dict[ProtoFile, tuple[list[str], list[str]]]):
        lines = ["export {"]
        if self.opts.enable_arena_switch:
            lines.append("  setArenaEnabled,")
        if self.opts.enable_json:
            lines.append("  toJsonString,")
        for proto_file, symbols in symbol_map.items():
            lines.append(f"  // {proto_file.fd.name}")
            lines.extend([f"  {enum}," for enum in symbols[0]])
            lines.extend([f"  {clazz}, ${clazz}," for clazz in symbols[1]])
        lines.append("} from 'MODIFY_ME_TO_YOUR_REAL_SO_NAME';")
        return lines
